import cv2 as cv
import numpy as np
import math as m
from matplotlib import pyplot as plt
from PyQt5.QtWidgets import *
from .Linear_UI import Ui_Dialog as UI_Linear
from .Log_UI import Ui_Dialog as UI_Log
from .Exp_UI import Ui_Dialog as UI_Exp
from .Pow_UI import Ui_Dialog as UI_Pow
from .Conv_UI import Ui_Dialog as UI_Conv
import time

def gray_deal(img, deal_Type):
    """根据用户的选择，对于图像做相应的灰度增强处理"""
    if img.shape[-1] == 3:
        pass
    if deal_Type == 1:
        img = linear_strench(img)
    elif deal_Type == 2:
        img = log_strench(img)
    elif deal_Type == 3:
        img = exp_strench(img)
    elif deal_Type == 4:
        img = pow_strench(img)
    return img

def linear_strench(img):
    """*功能 : 根据传入的图像及给定的c,d两个灰值区间参数值，进行线性拉伸
    *注意，只对灰度图像拉伸，函数：g(x,y)=(d-c)/(b-a)*[f(x,y)-a]+c=k*[f(x,y)-a]+c"""
    rows, cols = img.shape[:2]  # 获取宽和高
    new_img = np.zeros(img.shape, dtype=np.uint8)  # 新建同原图大小一致的空图像
    a, b = np.min(img), np.max(img)
    q_dialog = QDialog()
    dialog = UI_Linear()
    dialog.setupUi(q_dialog) # 继承QDialog()， 使得dialog具有show()方法
    dialog.lineEdit_a.setText(str(a))  # 显示原图灰度范围
    dialog.lineEdit_b.setText(str(b))
    dialog.lineEdit_c.setText(str(a))  # 初始化变换后灰度范围
    dialog.lineEdit_d.setText(str(b))
    q_dialog.show()
    if q_dialog.exec() == QDialog.Accepted:  # 提取用户交互变换后灰度范围
        c = int(dialog.lineEdit_c.text())
        d = int(dialog.lineEdit_d.text())
        time1 = time.time() # 程序计时开始

       # 这里写入你的核心代码
        for i in range(rows):
           for j in range(cols):
               new_img[i, j] = ((d-c)/(b-a))*(img[i,j]-a) + c

        time2 = time.time() # 程序计时结束
        hand_time = (time2 - time1) * 1000
        print("手写线性拉伸程序处理时间：%.3f毫秒" % hand_time)

        # OpenCV实现
        time1 = time.time()
        opencv_img = cv.normalize(img, None, c, d, cv.NORM_MINMAX)
        time2 = time.time()
        opencv_time = (time2 - time1) * 1000
        print("OpenCV线性拉伸程序处理时间：%.3f毫秒" % opencv_time)
        print("OpenCV比手写快：%.3f毫秒，加速比：%.2f倍" % (hand_time - opencv_time,
                                                         hand_time / opencv_time if opencv_time > 0 else float('inf')))
        return new_img

def log_strench(img):
    """*功能 : 根根据传入的图像及给定的a,b,c三个参数值，进行对数非线性拉伸
    *注意，只对灰度进行拉伸，函数：g(x,y)=a+lg[f(x,y)+1]/(c*lgb)"""
    rows, cols = img.shape[:2]  # 获取宽和高
    new_img = np.zeros(img.shape, dtype=np.uint8)  # 新建同原图大小一致的空图像
    q_dialog = QDialog()
    dialog = UI_Log()
    dialog.setupUi(q_dialog) # 继承QDialog()， 使得dialog具有show()方法
    dialog.lineEdit_a.setText(str(0.0))  # 初始化对数变换参数
    dialog.lineEdit_b.setText(str(2.0))
    dialog.lineEdit_c.setText(str(0.03))
    q_dialog.show()
    if q_dialog.exec() == QDialog.Accepted:  # 提取用户交互的参数
        a, b, c = float(dialog.lineEdit_a.text()), float(dialog.lineEdit_b.text()), float(dialog.lineEdit_c.text())
        if c == 0 or b <= 0 or b == 1: # 对参数进行预判断
            if QMessageBox(QMessageBox.Warning, '警告', '参数设置不合理！').exec():
                return img

        # 手写实现
        time1 = time.time()  # 程序计时开始
        img_float = img.astype(np.float64)
        transformed_values = a + np.log10(img_float + 1) / (c * np.log10(b))
        new_img = np.clip(transformed_values, 0, 255).astype(np.uint8)
        time2 = time.time()  # 程序计时结束
        hand_time = (time2 - time1) * 1000
        print("手写对数变换程序处理时间：%.3f毫秒" % hand_time)

        # OpenCV实现
        time1 = time.time()
        img_float = img.astype(np.float64)
        opencv_img = a + cv.log(img_float + 1) / (c * np.log(b))
        opencv_img = cv.normalize(opencv_img, None, 0, 255, cv.NORM_MINMAX)
        opencv_img = np.clip(opencv_img, 0, 255).astype(np.uint8)
        time2 = time.time()
        opencv_time = (time2 - time1) * 1000
        print("OpenCV对数变换程序处理时间：%.3f毫秒" % opencv_time)
        print("OpenCV比手写快：%.3f毫秒，加速比：%.2f倍" % (hand_time - opencv_time,
                                                         hand_time / opencv_time if opencv_time > 0 else float(
                                                             'inf')))
        return new_img

def exp_strench(img):
    """*功能 : 根根据传入的图像及给定的a,b,c三个参数值，进行指数非线性拉伸
    *注意，只对灰度进行拉伸，函数：g(x,y)=b^c[f(x,y)-a]-1"""
    rows, cols = img.shape[:2]  # 获取宽和高
    new_img = np.zeros(img.shape, dtype=np.uint8)  # 新建同原图大小一致的空图像
    q_dialog = QDialog()
    dialog = UI_Exp()
    dialog.setupUi(q_dialog) # 继承QDialog()， 使得dialog具有show()方法
    dialog.lineEdit_a.setText(str(150))  # 初始化对数变换参数
    dialog.lineEdit_b.setText(str(1.5))
    dialog.lineEdit_c.setText(str(0.6))
    q_dialog.show()
    if q_dialog.exec() == QDialog.Accepted:  # 提取用户交互的参数
        a, b, c = float(dialog.lineEdit_a.text()), float(dialog.lineEdit_b.text()), float(dialog.lineEdit_c.text())
        if b <= 0 or b == 1: # 对参数进行预判断
            if QMessageBox(QMessageBox.Warning, '警告', '参数设置不合理！').exec():
                return img

        # 这里写入你的核心代码
        # 手写实现
        time1 = time.time()  # 程序计时开始
        img_float = img.astype(np.float64)
        transformed_values = np.power(b, c * (img_float - a)) - 1
        new_img = np.clip(transformed_values, 0, 255).astype(np.uint8)
        time2 = time.time()  # 程序计时结束
        hand_time = (time2 - time1) * 1000
        print("手写指数变换程序处理时间：%.3f毫秒" % hand_time)

        # OpenCV实现
        time1 = time.time()
        img_float = img.astype(np.float64)
        opencv_img = cv.exp(c * (img_float - a) * np.log(b)) - 1
        opencv_img = np.clip(opencv_img, 0, 255).astype(np.uint8)
        time2 = time.time()
        opencv_time = (time2 - time1) * 1000
        print("OpenCV指数变换程序处理时间：%.3f毫秒" % opencv_time)
        print("OpenCV比手写快：%.3f毫秒，加速比：%.2f倍" % (hand_time - opencv_time,
                                                         hand_time / opencv_time if opencv_time > 0 else float('inf')))

    return new_img

def pow_strench(img):
    """*功能 : 根根据传入的图像及给定的c,r两个参数值，进行幂律非线性拉伸
    *注意，只对灰度进行拉伸，函数：g(x,y)=c[f(x,y)]^r"""
    rows, cols = img.shape[:2]  # 获取宽和高
    new_img = np.zeros(img.shape, dtype=np.uint8)  # 新建同原图大小一致的空图像
    q_dialog = QDialog()
    dialog = UI_Pow()
    dialog.setupUi(q_dialog) # 继承QDialog()， 使得dialog具有show()方法
    dialog.lineEdit_c.setText(str(1))  # 初始化对数变换参数
    dialog.lineEdit_r.setText(str(1.5))
    q_dialog.show()
    if q_dialog.exec() == QDialog.Accepted:  # 提取用户交互的参数
        c, r = float(dialog.lineEdit_c.text()), float(dialog.lineEdit_r.text())
        if r <= 0 or c <= 0: # 对参数进行预判断
            if QMessageBox(QMessageBox.Warning, '警告', '参数设置不合理！').exec():
                return img

        time1 = time.time() # 程序计时开始

        # 这里写入你的核心代码
        # 手写实现
        time1 = time.time()  # 程序计时开始
        img_float = img.astype(np.float64)
        transformed_values = c * np.power(img_float, r)
        min_val = np.min(transformed_values)
        max_val = np.max(transformed_values)
        normalized_values = (transformed_values - min_val) / (max_val - min_val) * 255.0
        new_img = np.clip(normalized_values, 0, 255).astype(np.uint8)
        time2 = time.time()  # 程序计时结束
        hand_time = (time2 - time1) * 1000
        print("手写幂律变换程序处理时间：%.3f毫秒" % hand_time)

        # OpenCV实现（使用伽马校正）
        time1 = time.time()
        img_normalized = img.astype(np.float32) / 255.0
        opencv_img = c * np.power(img_normalized, r)
        opencv_img = np.clip(opencv_img * 255, 0, 255).astype(np.uint8)
        time2 = time.time()
        opencv_time = (time2 - time1) * 1000
        print("OpenCV幂律变换程序处理时间：%.3f毫秒" % opencv_time)
        print("OpenCV比手写快：%.3f毫秒，加速比：%.2f倍" % (hand_time - opencv_time,
                                                         hand_time / opencv_time if opencv_time > 0 else float('inf')))

    return new_img


def hist_equalization(img, jug):
    """*功能 : 直方图均衡化算法, jug判断返回是图像/直方图"""
    rows, cols = img.shape[:2]  # 获取宽和高
    new_img = np.zeros(img.shape, dtype=np.uint8)  # 新建同原图大小一致的空图像
    total_pixels = rows * cols
    time1 = time.time()  # 程序计时开始
    hist = creat_histogram(img)

    # 直方图均衡化核心代码
    if img.ndim == 2:  # 灰度图像
        # 计算累积分布函数
        cdf = [0] * 256
        cdf[0] = hist[0]
        for i in range(1, 256):
            cdf[i] = cdf[i - 1] + hist[i]

        # 归一化累积分布函数并映射到0-255范围
        cdf_min = min([val for val in cdf if val > 0])  # 找到第一个非零累积值
        for i in range(256):
            if cdf[i] > 0:
                cdf[i] = round((cdf[i] - cdf_min) / (total_pixels - cdf_min) * 255)

        # 应用映射到新图像
        for row in range(rows):
            for col in range(cols):
                new_img[row][col] = cdf[img[row][col]]

    elif img.ndim == 3:  # 彩色图像
        # 对每个颜色通道分别进行均衡化
        for channel in range(3):
            # 计算当前通道的累积分布函数
            cdf = [0] * 256
            cdf[0] = hist[channel][0]
            for i in range(1, 256):
                cdf[i] = cdf[i - 1] + hist[channel][i]

            # 归一化累积分布函数并映射到0-255范围
            cdf_min = min([val for val in cdf if val > 0])  # 找到第一个非零累积值
            for i in range(256):
                if cdf[i] > 0:
                    cdf[i] = round((cdf[i] - cdf_min) / (total_pixels - cdf_min) * 255)

            # 应用映射到新图像的当前通道
            for row in range(rows):
                for col in range(cols):
                    new_img[row][col][channel] = cdf[img[row][col][channel]]

    time2 = time.time()  # 程序计时结束
    print(f"直方图均衡化耗时: {time2 - time1:.4f}秒")


    if jug:
        imgs = [img, new_img]
        colors = ("b", "r")
        texts = ("original histogram", "histogram after equalization")
        for i in range(2):
            hist = cv.calcHist([imgs[i]], [0], None, [256], [0, 255])
            plt.plot(hist, color=colors[i], label=texts[i])
        plt.xlim([0, 256])
        plt.legend()
        plt.show()

    return new_img
def creat_histogram(img):
    """*功能 : 计算传入图像直方图，若是彩色图像，计算各颜色分量直方图并返回"""
    rows, cols = img.shape[:2]  # 获取宽和高
    hist = []
    if img.ndim == 2: # 灰度图像统计直方图
        hist = [0] * 256  # 建立灰度图像直方图
        # 图像遍历
        for row in range(rows):
            for col in range(cols):
                hist[img[row][col]] += 1
    elif img.ndim == 3:  # 彩色图像统计直方图
        hist = [[0] * 256, [0] * 256, [0] * 256]  # 建立彩色图像直方图
        # 图像遍历
        for row in range(rows):
            for col in range(cols):
                hist[0][img[row][col][0]] += 1
                hist[1][img[row][col][1]] += 1
                hist[2][img[row][col][2]] += 1
    return hist

def gray_smooth(img, deal_Type):
    """根据用户的选择，对于图像做相应的图像平滑处理"""
    if img.shape[-1] == 3:
        pass
    if deal_Type == 1:
        img = neighbor_average(img)
    elif deal_Type == 2:
        img = median_filter(img)
    return img

def neighbor_average(img):
    """*功能 : 用户交互卷积模板，获取卷积系数进行邻域平滑，只对灰度图像处理"""
    rows, cols = img.shape[:2]  # 获取宽和高
    new_img = np.zeros(img.shape, dtype=np.uint8)  # 新建同原图大小一致的空图像
    q_dialog = QDialog()
    dialog = UI_Conv()
    dialog.setupUi(q_dialog) # 继承QDialog()， 使得dialog具有show()方法
    q_dialog.show()
    if q_dialog.exec() == QDialog.Accepted:  # 提取用户交互的参数
        np_kernel = np.array([[float(dialog.lineEdit1.text()), float(dialog.lineEdit2.text()), float(dialog.lineEdit3.text())],
                             [float(dialog.lineEdit4.text()), float(dialog.lineEdit5.text()), float(dialog.lineEdit6.text())],
                             [float(dialog.lineEdit7.text()), float(dialog.lineEdit8.text()), float(dialog.lineEdit9.text())]])
        np_kernel = np_kernel/np_kernel.sum() # 正则化

        time1 = time.time() # 程序计时开始

        # 这里写入你的核心代码
        # 手写实现
        time1 = time.time()  # 程序计时开始
        for i in range(1, rows - 1):
            for j in range(1, cols - 1):
                # 使用numpy的element-wise乘法
                neighborhood = img[i - 1:i + 2, j - 1:j + 2]
                result = np.sum(neighborhood * np_kernel)
                new_img[i, j] = np.clip(result, 0, 255).astype(np.uint8)

            # 复制边界
        new_img[0, :] = img[0, :]
        new_img[rows - 1, :] = img[rows - 1, :]
        new_img[:, 0] = img[:, 0]
        new_img[:, cols - 1] = img[:, cols - 1]
        time2 = time.time()  # 程序计时结束
        hand_time = (time2 - time1) * 1000
        print("手写邻域平均平滑程序处理时间：%.3f毫秒" % hand_time)

        # OpenCV实现
        time1 = time.time()
        opencv_img = cv.filter2D(img, -1, np_kernel, borderType=cv.BORDER_REPLICATE)
        time2 = time.time()
        opencv_time = (time2 - time1) * 1000
        print("OpenCV邻域平均平滑程序处理时间：%.3f毫秒" % opencv_time)
        print("OpenCV比手写快：%.3f毫秒，加速比：%.2f倍" % (hand_time - opencv_time,
                                                         hand_time / opencv_time if opencv_time > 0 else float('inf')))

    return new_img


def median_filter(img):
    """*功能 : 中值滤波"""
    rows, cols = img.shape[:2]  # 获取宽和高
    new_img = np.zeros(img.shape, dtype=np.uint8)  # 新建同原图大小一致的空图像
    len = 3  # 定义中值滤波模板3×3

    time1 = time.time() # 程序计时开始

    # 这里写入你的核心代码
    # 手写实现
    time1 = time.time()  # 程序计时开始
    # 遍历图像，跳过边界像素
    for i in range(1, rows - 1):
        for j in range(1, cols - 1):
            # 1. 提取当前像素的3x3邻域
            neighborhood = img[i - 1:i + 2, j - 1:j + 2]

            # 2. 将邻域数组展平为一维数组
            flattened_neighborhood = neighborhood.flatten()

            # 3. 对一维数组进行排序
            sorted_neighborhood = np.sort(flattened_neighborhood)

            # 4. 找到中值。对于3x3的模板，索引为4的元素是第5个值，即中值。
            median_value = sorted_neighborhood[4]  # 索引从0开始，[0,1,2,3,4,5,6,7,8]

            # 5. 将中值赋给新图像的对应位置
            new_img[i, j] = median_value

    # 复制边界像素
    new_img[0, :] = img[0, :]
    new_img[rows - 1, :] = img[rows - 1, :]
    new_img[:, 0] = img[:, 0]
    new_img[:, cols - 1] = img[:, cols - 1]
    time2 = time.time()  # 程序计时结束
    hand_time = (time2 - time1) * 1000
    print("手写中值滤波程序处理时间：%.3f毫秒" % hand_time)

    # OpenCV实现
    time1 = time.time()
    opencv_img = cv.medianBlur(img, 3)
    time2 = time.time()
    opencv_time = (time2 - time1) * 1000
    print("OpenCV中值滤波程序处理时间：%.3f毫秒" % opencv_time)
    print("OpenCV比手写快：%.3f毫秒，加速比：%.2f倍" % (hand_time - opencv_time,
                                                     hand_time / opencv_time if opencv_time > 0 else float('inf')))

    return new_img
